#include "optimizer.h"

Optimizer::Optimizer(std::shared_ptr<ImuPreIntegration> imu_preintegration, std::shared_ptr<WheelPreIntegration> wheel_preintegration) {
    imu_factor_ = std::make_shared<ImuFactor>(imu_preintegration);
    wheel_factor_ = std::make_shared<WheelFactor>(wheel_preintegration);;
}
// ti, vxi, vyi, pxi, pyi, bgi, baxi, bayi, tj, vxj, vyj, pxj, pyj
void Optimizer::OptimizeLM(State& last_state, State& state, const Eigen::Vector3d &last_velocity, const Eigen::Vector3d &velocity, const Eigen::Vector3d &last_pose, const Eigen::Vector3d &curr_pose) {
    float cost_old = 0.0f;
    bool recompute = true;
    double eps = 10000;
    int runs = 0;
    double lamda =  1e-3;
    double lam_mult = 10;
    double min_dcost = 1e-3;

    State last_imu_state = last_state; 
    State imu_state = state;

    Eigen::Matrix<double, 13, 13> Hess = Eigen::Matrix<double, 13, 13>::Zero();
    Eigen::Matrix<double, 13, 1> grad = Eigen::Matrix<double, 13, 1>::Zero();

    Eigen::Matrix<double, 5, 1> imufacotr_residual = Eigen::Matrix<double, 5, 1>::Zero();

    Eigen::Matrix<double, 5, 5> imufacotr_information_matrix = Eigen::Matrix<double, 5, 5>::Identity();
    imufacotr_information_matrix.block<5,5>(0,0) = imu_factor_->GetCov().inverse();
    
    // cout << "imufacotr_information_matrix =========== : \n" << imufacotr_information_matrix << endl;

    Eigen::Matrix<double, 5, 8> imufacotr_jacobianXi = Eigen::Matrix<double, 5, 8>::Zero();
    Eigen::Matrix<double, 5, 5> imufacotr_jacobianXj = Eigen::Matrix<double, 5, 5>::Zero();

    imu_factor_->ComputeResidualAndJacobianPoseGraph(last_state, state, imufacotr_residual, imufacotr_jacobianXi, imufacotr_jacobianXj);
    cost_old += Chi2(imufacotr_residual, imufacotr_information_matrix);

    // get odometry preintegration residual and jacobian
    Eigen::Vector3d wheel_preint_residual;
    Eigen::Matrix<double, 3, 8> wheelfacotr_jacobianXi;
    Eigen::Matrix<double, 3, 5> wheelfacotr_jacobianXj;

    wheel_factor_->ComputeResidualAndJacobianPoseGraph(last_state, state, wheel_preint_residual, wheelfacotr_jacobianXi, wheelfacotr_jacobianXj);
    Eigen::Matrix<double, 3, 3> wheelfacotr_information_matrix = Eigen::Matrix<double, 3, 3>::Identity();
    wheelfacotr_information_matrix.block<3,3>(0,0) = wheel_factor_->GetCov().inverse();
    cost_old += Chi2(wheel_preint_residual, wheelafcotr_information_matrix);

    double last_pose_sigma = 1e-15;
    double curr_pose_sigma = 1e-15;
    Eigen::Matrix<double, 5, 5> last_pose_covariance = Eigen::Matrix<double, 5, 5>::Identity();
    last_pose_covariance(0,0) = std::pow(last_pose_sigma,2);
    last_pose_covariance(1,1) = std::pow(last_pose_sigma,2);
    last_pose_covariance(2,2) = std::pow(last_pose_sigma,2);
    last_pose_covariance(3,3) = std::pow(last_pose_sigma,2);
    last_pose_covariance(4,4) = std::pow(last_pose_sigma,2);
    Eigen::Matrix<double, 5, 5> last_information_matrix = last_pose_covariance.inverse();

    Eigen::Matrix<double, 5, 5> curr_pose_covariance = Eigen::Matrix<double, 5, 5>::Identity();
    curr_pose_covariance(0,0) = std::pow(curr_pose_sigma,2);
    curr_pose_covariance(1,1) = std::pow(curr_pose_sigma,2);
    curr_pose_covariance(2,2) = std::pow(curr_pose_sigma,2);
    curr_pose_covariance(3,3) = std::pow(curr_pose_sigma,2);
    curr_pose_covariance(4,4) = std::pow(curr_pose_sigma,2);
    Eigen::Matrix<double, 5, 5> curr_information_matrix = curr_pose_covariance.inverse();
    
    // cout << "curr_information_matrix == : \n" << curr_information_matrix << endl;

    Eigen::Matrix<double, 10, 10> pose_information_matrix = Eigen::Matrix<double, 10, 10>::Identity();
    pose_information_matrix.block<5,5>(0,0) = last_information_matrix;
    pose_information_matrix.block<5,5>(5,5) = curr_information_matrix;
    Eigen::Matrix<double, 5, 1> residualXi; 
    Eigen::Matrix<double, 5, 8> jacobianXi;

    Eigen::Matrix<double, 5, 1> residualXj; 
    Eigen::Matrix<double, 5, 5> jacobianXj;

    Eigen::Matrix<double, 5, 5> jacobianXij;
    Eigen::Matrix<double, 5, 8> jacobianXji;
    PoseObserve::ComputeResidualAndJacobianPoseGraph(last_state, last_pose, last_velocity.head<2>(), state, curr_pose, velocity.head<2>(), residualXi, jacobianXi, residualXj, jacobianXj, jacobianXij, jacobianXji);
    cost_old += Chi2(residualXi, last_information_matrix);
    cost_old += Chi2(residualXj, curr_information_matrix);

    // cout << "wheelfacotr_information_matrix =========== : \n" << wheelfacotr_information_matrix << endl;

    // cout << "cost_old =========== : " << cost_old << endl;

    while (runs < 100 && lamda < 1e10 && eps > 1e-6) {
        
        if (recompute) {
            Hess.setZero();
            grad.setZero();

            imu_factor_->ComputeResidualAndJacobianPoseGraph(last_state, state, imufacotr_residual, imufacotr_jacobianXi, imufacotr_jacobianXj);
            // imufacotr_information_matrix.block<5,5>(0,0) = imu_factor_->GetCov().inverse();

            Eigen::Matrix<double,5,13> iH = Eigen::Matrix<double,5,13>::Zero();
            iH.block<5,8>(0,0) = imufacotr_jacobianXi;
            iH.block<5,5>(0,8) = imufacotr_jacobianXj;
    
            wheel_factor_->ComputeResidualAndJacobianPoseGraph(last_state, state, wheel_preint_residual, wheelfacotr_jacobianXi, wheelfacotr_jacobianXj);
            // wheelfacotr_information_matrix.block<3,3>(0,0) = wheel_factor_->GetCov().inverse();

            Eigen::Matrix<double,3,13> wH = Eigen::Matrix<double,3,13>::Zero();
            wH.block<3,8>(0,0) = wheelfacotr_jacobianXi;
            wH.block<3,5>(0,8) = wheelfacotr_jacobianXj;

            Eigen::Matrix<double,18,13> H = Eigen::Matrix<double,18,13>::Zero();
            H.block<5,13>(0,0) = iH;
            H.block<3,13>(5,0) = wH;
            H.block<5,8>(8,0) = jacobianXi;
            H.block<5,5>(13,8) = jacobianXj;
            Eigen::Matrix<double, 18, 18> information_matrix = Eigen::Matrix<double, 18, 18>::Identity();
            information_matrix.block<5,5>(0,0) = imufacotr_information_matrix;
            information_matrix.block<3,3>(5,5) = wheelfacotr_information_matrix;
            information_matrix.block<10,10>(8,8) = pose_information_matrix;
            Eigen::Matrix<double, 18, 1> residual;
            residual.block<5,1>(0,0) = imufacotr_residual;
            residual.block<3,1>(5,0) = wheel_preint_residual;
            residual.block<5,1>(8,0) = residualXi;
            residual.block<5,1>(13,0) = residualXj;
            grad.noalias() += -H.transpose() * information_matrix * residual;
            Hess.noalias() += H.transpose() * information_matrix * H;        
        }

        // Solve Levenberg iteration
        Eigen::Matrix<double,13, 13> Hess_l = Hess;
        for (size_t r=0; r < (size_t)Hess.rows(); r++) {
            Hess_l(r,r) *= (1.0 + lamda);
            //std::cout << " lamda : " << lamda << std::endl;
        }
        
        // std::cout << " lamda : " << lamda << std::endl;

        Eigen::Matrix<double,13,1> dx = Hess_l.colPivHouseholderQr().solve(grad);

        std::cout << " iterator idx and dx = : " << runs << ", " << dx.transpose() << std::endl;

        State last_state_tmp = last_state;
        State state_tmp = state;
        update_state(last_state_tmp, state_tmp, dx);
        imu_factor_->ComputeResidualAndJacobianPoseGraph(last_state_tmp, state_tmp, imufacotr_residual, imufacotr_jacobianXi, imufacotr_jacobianXj);
        wheel_factor_->ComputeResidualAndJacobianPoseGraph(last_state_tmp, state_tmp, wheel_preint_residual, wheelfacotr_jacobianXi, wheelfacotr_jacobianXj);
        
        PoseObserve::ComputeResidualAndJacobianPoseGraph(last_state_tmp, last_pose, last_velocity.head<2>(), state_tmp, curr_pose, velocity.head<2>(), residualXi, jacobianXi, residualXj, jacobianXj, jacobianXij, jacobianXji);

        // imufacotr_information_matrix.block<5,5>(0,0) = imu_factor_->GetCov().inverse();
        // wheelfacotr_information_matrix.block<3,3>(0,0) = wheel_factor_->GetCov().inverse();

        double cost = 0.0;
        cost += Chi2(imufacotr_residual, imufacotr_information_matrix);
        
        // std::cout << " imu cost : " << cost << std::endl;   

        cost += Chi2(wheel_preint_residual, wheelfacotr_information_matrix);
        
        cost += Chi2(residualXi, last_information_matrix);
        cost += Chi2(residualXj, curr_information_matrix);
        // std::cout << " wheel cost : " << Chi2(wheel_preint_residual, wheelfacotr_information_matrix) << std::endl;   

        // std::cout << " cost_old : " << cost_old << std::endl;
        // std::cout << " cost : " << cost << std::endl;   
        // std::cout << " (cost_old - cost) / cost_old : " << (cost_old - cost) / cost_old << std::endl;   

        // Check if converged
        if (cost <= cost_old && (cost_old - cost) / cost_old < min_dcost) {
            eps = 0;

            update_state(last_state, state, dx);

            // std::cout << "converged" << std::endl;
            // std::cout << " cost : " << cost << std::endl;
            // std::cout << " cost_old : " << cost_old << std::endl;
            // std::cout << " runs : " << runs << std::endl;
            // std::cout << " dx ===== : " << dx.transpose() << std::endl;
            break;
        }

        // If cost is lowered, accept step
        if (cost <= cost_old) {
            // std::cout << " cost_old : " << cost_old << std::endl;
            // std::cout << " cost : " << cost << std::endl;                
            // std::cout << " runs : " << runs << std::endl;

            recompute = true;
            cost_old = cost;

            update_state(last_state, state, dx);

            runs++;
            lamda = lamda / lam_mult;
            eps = dx.norm();

        } else {
            recompute = false;
            lamda = lamda * lam_mult;
            continue;
        }      
    }
}

double Optimizer::Chi2(const Eigen::Matrix<double, 5, 1>& residual, const Eigen::Matrix<double, 5, 5>& information_matrix) {
	return residual.transpose() * information_matrix * residual;
}

double Optimizer::Chi2(const Eigen::Matrix<double, 2, 1>& residual, const Eigen::Matrix<double, 2, 2>& information_matrix) {
    return residual.transpose() * information_matrix * residual;
}

double Optimizer::Chi2(const Eigen::Matrix<double, 3, 1>& residual, const Eigen::Matrix<double, 3, 3>& information_matrix) {
    return residual.transpose() * information_matrix * residual;
}

void Optimizer::update_state(State& last_state, State& state, const Eigen::Matrix<double, 13, 1>& dx) {
    last_state.R = last_state.R * SO2(dx[0]);
    last_state.v += dx.block<2,1>(1,0);
    last_state.p += dx.block<2,1>(3,0);
    last_state.bg += dx(5,0);
    last_state.ba += dx.block<2,1>(6,0); 
    state.R = state.R * SO2(dx[8]);
    state.v += dx.block<2,1>(9,0);
    state.p += dx.block<2,1>(11,0);
    state.bg = last_state.bg; 
    state.ba = last_state.ba;
}
